"""
递归进行文本分块，分隔符: ['\n\n', '\n', '。', '；', ';']
"""
from collections.abc import Iterable
from typing import Optional, List, Dict
from core.document import Document
import copy
import re
from loguru import logger


class TextSplitter(object):
    def __init__(self,
        chunk_size: int = 300,
        chunk_overlap: int = 50,
        keep_separator: bool = True,
        add_start_index: bool = True
    ) -> None:
        """ Args:
            - chunk_size: 片段的最大字数
            - chunk_overlap: 相邻片段的重叠字数
            - keep_separator: 是否保留片段的分割符
            - add_start_index: 是否在metadata中加入片段的起始位置
        """
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap
        self._keep_separator = keep_separator
        self._add_start_index = add_start_index
        self._separators = ['\n\n', '\n', '。', '；', ';']
        self._length_function = lambda x: [len(x) for x in x]

    def create_documents(self, text: str, metadata: Optional[Dict] = None) -> List[Document]:
        """ 文本分割 """
        metadata = metadata or {}
        index = -1
        documents = []
        last_chunk = None
        for chunk in self.split_text(text, self._separators):
            copy_metadata = copy.deepcopy(metadata)
            if self._add_start_index:
                index = text.find(chunk, index + 1)
                copy_metadata["start_index"] = index
                # 如果段落文字过多，则我们只取段落的一部分用于召回。
                if index > self._chunk_size and last_chunk is not None:
                    copy_metadata['block_content'] = last_chunk + chunk
            new_doc = Document(content=chunk, metadata=copy_metadata)
            documents.append(new_doc)
            last_chunk = chunk
        return documents

    def _join_docs(self, docs: List[str], separator: str) -> Optional[str]:
        text = separator.join(docs)
        text = text.strip()
        if text == "":
            return None
        else:
            return text

    def _merge_splits(self, splits: Iterable[str], separator: str, lengths: List[int]) -> List[str]:
        # We now want to combine these smaller pieces into medium size
        # chunks to send to the LLM.
        separator_len = self._length_function([separator])[0]

        docs = []
        current_doc: List[str] = []
        total = 0
        index = 0
        for d in splits:
            _len = lengths[index]
            if total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size:
                if total > self._chunk_size:
                    logger.warning(
                        f"Created a chunk of size {total}, which is longer than the specified {self._chunk_size}"
                    )
                if len(current_doc) > 0:
                    doc = self._join_docs(current_doc, separator)
                    if doc is not None:
                        docs.append(doc)
                    # Keep on popping if:
                    # - we have a larger chunk than in the chunk overlap
                    # - or if we still have any chunks and the length is long
                    while total > self._chunk_overlap or (
                        total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size and total > 0
                    ):
                        total -= self._length_function([current_doc[0]])[0] + (
                            separator_len if len(current_doc) > 1 else 0
                        )
                        current_doc = current_doc[1:]
            current_doc.append(d)
            total += _len + (separator_len if len(current_doc) > 1 else 0)
            index += 1
        doc = self._join_docs(current_doc, separator)
        if doc is not None:
            docs.append(doc)
        return docs

    def split_text(self, text: str, separators: List[str]) -> List[str]:
        final_chunks = []
        separator = separators[-1]
        new_separators = []

        for i, _s in enumerate(separators):
            if _s == "":
                separator = _s
                break
            if re.search(_s, text):
                separator = _s
                new_separators = separators[i + 1 :]
                break

        splits = self._split_text_with_regex(text, separator, self._keep_separator)
        _good_splits = []
        _good_splits_lengths = []  # cache the lengths of the splits
        _separator = "" if self._keep_separator else separator
        s_lens = self._length_function(splits)
        for s, s_len in zip(splits, s_lens):
            if s_len < self._chunk_size:
                _good_splits.append(s)
                _good_splits_lengths.append(s_len)
            else:
                if _good_splits:
                    merged_text = self._merge_splits(_good_splits, _separator, _good_splits_lengths)
                    final_chunks.extend(merged_text)
                    _good_splits = []
                    _good_splits_lengths = []
                if not new_separators:
                    final_chunks.append(s)
                else:
                    other_info = self.split_text(s, new_separators)
                    final_chunks.extend(other_info)

        if _good_splits:
            merged_text = self._merge_splits(_good_splits, _separator, _good_splits_lengths)
            final_chunks.extend(merged_text)

        return final_chunks

    @staticmethod
    def _split_text_with_regex(text: str, separator: str, keep_separator: bool) -> List[str]:
        # Now that we have the separator, split the text
        if separator:
            if keep_separator:
                # The parentheses in the pattern keep the delimiters in the result.
                _splits = re.split(f"({re.escape(separator)})", text)
                splits = [_splits[i - 1] + _splits[i] for i in range(1, len(_splits), 2)]
                if len(_splits) % 2 != 0:
                    splits += _splits[-1:]
            else:
                splits = re.split(separator, text)
        else:
            splits = list(text)
        return [s for s in splits if (s not in {"", "\n"})]


if __name__ == '__main__':

    text = """
    根据《中华人民共和国民法典》及相关法律法规的规定，甲、乙双方在平等、自愿的基础上，就甲方将柳州市南环路北侧两幅地块出租给乙方的相关事宜，自愿订立本合同，以明确双方权利义务，共同信守履行。
    第一条  出租标的基本情况
    建设用地使用权人为甲方，承租地块信息详见附件，土地使用权未设定抵押。乙方对租赁土地的现有情况（土地性质、现状）已有充分的认知。
    第二条  承租地块范围内的建（构）筑物情况
    A1地块有一栋占地约1,500平方米三层高混凝土框架结构建筑，原建设作为售房部使用（以下简称原售房部），本合同签订时未装修完毕。承租地块租赁期间，乙方可以使用原售房部用于办公。乙方因使用原售房部所产生的装修、采购办公用品等费用由乙方自行支付，不得以任何原因要求甲方支付或偿还。
    第三条  租赁面积
    乙方承租地块面积：柳州市柳南区南环路北侧A1、A2地块：37842.76平方米（折合56.76亩）；柳州市柳南区南环路北侧B1、B2、B3地块：85,249.14平方米（合127.87亩）。
    第四条  租赁期限
    为便于乙方使用承租地块，本合同签订后，甲乙双方共同完成对承租地块上原承租户的清退工作，甲方根据乙方用地需求，配合乙方完成清退原租户、土地平整等工作。当承租地块达到以下两个条件之一的视为达到交付条件，双方应签署土地交付单，正式交付土地，签署土地交付单之日起开始计算租期，租赁期限合计五年：
    （一）地块连续完整清退面积达到承租面积95%，剩余5%不能含有南环路沿街往里延伸30米部分。
    （二）乙方自行招商引入的次承租户在承租地块开始经营、乙方已利用承租地块产生了实际效益。
    """
    spliter = TextSplitter(chunk_size=300, chunk_overlap=50)
    docs = spliter.create_documents(text, metadata={})
    for doc in docs:
        print(doc)